package tree.huffman;

import java.io.Serializable;
import java.util.*;

/**
 * @Author Jun jie Yi
 * @Date 2021/7/21-16:24
 */
public class HuffmanCode {

    // 存储 huffman 编码表
    static Map<Byte, String> huffmanCodes = new HashMap<>();

    public static void main(String[] args) {
        String context = "i like like like java do you like a java";

        Map<Byte, String> codes = getCodes(context.getBytes());//获取huffman编码表
        byte[] bytes = huffmanZip(context.getBytes());//压缩后的huffman编码数组

        byte[] decode = decode(codes, bytes);//解压后的huffman数组

        System.out.println(Arrays.toString(decode));
        System.out.println(Arrays.toString(bytes));
        System.out.println(Arrays.toString(context.getBytes()));
    }

    /**
     * 1.先将数据装转化为二进制的字符串
     * 2.用这个字符串去对照huffman编码重新转成元数据
     * 完成对压缩字符串的解码
     *
     * @param huffmanCodes huffman编码表map
     * @param huffmanBytes 压缩后的数据
     * @return 原来的字符串编码数组
     */
    public static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {
        //1.先得到 huffmanBytes 的二进制字符串
        StringBuilder stringBuilder = new StringBuilder();
        //2.将byte数组转化为二进制的字符串
        for (int i = 0; i < huffmanBytes.length; i++) {
            //判断是否是最后一个字节
            boolean flag = (i == huffmanBytes.length - 1);
            stringBuilder.append(byteToBitString(!flag, huffmanBytes[i]));
        }
        //把字符串按照指定的huffman表进行解码
        //将huffmanCodes进行反转
        Map<String, Byte> map = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
            map.put(entry.getValue(), entry.getKey());
        }
        //创建一个集合存放 byte
        List<Byte> list = new ArrayList<>();
        //i可以理解为一个索引
        int i = 0;
        while(i < stringBuilder.length()) {
            int count = 1;//小的计数器
            boolean flag = true;
            Byte b = null;
            while(flag) {
                //递增的取出 key1
                //i 不动，让 count 移动，指定匹配到一个字符
                String key = stringBuilder.substring(i, i + count);
                b = map.get(key);
                if(b == null) {
                    count ++; //计数器 +1 往后面找进行匹配
                } else {
                    //匹配到
                    flag = false;
                }
            }
            list.add(b);
            i += count;//i 直接移动到 count
        }
        //将数据存入 b 数组
        byte[] b = new byte[list.size()];
        for (int k = 0; k < b.length; k++) {
            b[k] = list.get(k);
        }
        return b;
    }

    //将 byte 数据转化为字符串
    public static String byteToBitString(boolean flag, byte b) {
        //使用变量保存 b，将b转成int
        int temp = b;
        //这里如果是正数还需要补位 EG: 1 => 00000001
        if(flag) {
            temp |= 256;
        }
        //返回的是 temp 对应的二进制补码(注意是 int 的4字节补码)
        String str = Integer.toBinaryString(temp);
        if(flag) {
            return str.substring(str.length() - 8);//只取后8位
        } else {
            return str;//如果刚好是8位就不用补高位
        }
    }

    //获取根据huffman编码表压缩后的数据(对字符串进行压缩)
    public static byte[] huffmanZip(byte[] bytes) {
        List<Node> nodes = getNodes(bytes);//根据byte数组获取对应的集合
        Node huffmanTree = createHuffmanTree(nodes);//创建huffman树
        Map<Byte, String> codes = getCodes(huffmanTree);//获取huffman编码
        //根据huffman编码表压缩
        return zip(bytes, codes);//根据byte数组和huffman编码对数据进行压缩
    }

    //将文件对应的byte[]数组，通过生成的huffman编码表，返回一个huffman编码压缩后的byte[]
    public static byte[] zip(byte[] bytes, Map<Byte, String> huffmanCodes) {
        //1.利用huffman编码表，将bytes转化为编码表对应的字符串
        StringBuilder stringBuilder = new StringBuilder();
        for (byte b : bytes) {
            stringBuilder.append(huffmanCodes.get(b));
        }
        //获取长度
        int len = (stringBuilder.length() + 7) / 8;
        //创建一个编码后的byte数组
        byte[] huffmanCodeBytes = new byte[len];
        int index = 0;//记录是第几个byte
        //创建有压缩后的byte[]数组
        for (int i = 0; i < stringBuilder.length(); i = i + 8) {

            String strByte;
            if(i + 8 > stringBuilder.length()) {
                strByte = stringBuilder.substring(i);//最后一个不够8位，就一起放进去
            } else {
                strByte = stringBuilder.substring(i, i + 8);//够8位就放入8个
            }
            //将strByte转成一个byte
            //Integer.parseInt(strByte, 2) 将2进制的数转化为常用的10进制
            huffmanCodeBytes[index] = (byte) Integer.parseInt(strByte, 2);
            index ++;
        }
        return huffmanCodeBytes;
    }

    //根据 byte[] 数组获取编码表
    public static Map<Byte, String> getCodes(byte[] bytes) {
        List<Node> nodes = getNodes(bytes);
        Node huffmanTree = createHuffmanTree(nodes);
        return getCodes(huffmanTree);
    }

    //根据huffman树获取huffman编码表
    public static Map<Byte, String> getCodes(Node root) {
        if(root == null) {
            return null;
        }
        StringBuilder stringBuilder = new StringBuilder();
        getCodes(root.left, "0", stringBuilder);
        getCodes(root.right, "1", stringBuilder);
        return huffmanCodes;
    }

    /**
     * 功能：将传入的node节点的所有叶子节点的huffman编码得到
     * @param node
     * @param code
     * @param stringBuilder
     */
    public static void getCodes(Node node, String code, StringBuilder stringBuilder) {
        StringBuilder sb = new StringBuilder(stringBuilder);
        //将code加入到sb
        sb.append(code);
        //如果 node == null 不处理
        if(node != null && node.data == null) {
            //递归处理
            getCodes(node.left, "0", sb); //向左递归
            getCodes(node.right, "1", sb); //向右递归
        } else {
            //是一个叶子节点
            huffmanCodes.put(node.data, sb.toString());
        }
    }

    //创建Huffman树
    public static Node createHuffmanTree(List<Node> nodes) {
        while(nodes.size() > 1) {
            //1.先对其进行排序
            Collections.sort(nodes);
            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);
            Node parent = new HuffmanCode().new Node(null, leftNode.weight + rightNode.weight);
            parent.left = leftNode;
            parent.right = rightNode;
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            nodes.add(parent);
        }
        return nodes.get(0);
    }

    //统计每个字符出现的次数
    public static List<Node> getNodes(byte[] bytes) {
        //创建一个ArrayList
        List<Node> nodes = new ArrayList<>();
        //遍历nodes统计每一个字符出现的次数
        Map<Byte, Integer> counts = new HashMap<>();
        for (byte b : bytes) {
            Integer count = counts.get(b);
            if(count == null) {
                counts.put(b, 1);
            } else {
                counts.put(b, ++count);
            }
        }
        for (Map.Entry<Byte, Integer> entry : counts.entrySet()) {
            nodes.add(new HuffmanCode().new Node(entry.getKey(), entry.getValue()));
        }
        return nodes;
    }

    //实现前序遍历
    public static void preOrder(Node root) {
        if(root == null) return;
        System.out.println(root);
        if(root.left != null) preOrder(root.left);
        if(root.right != null) preOrder(root.right);
    }

    //创建结点的类
    class Node implements Comparable<Node>, Serializable {
        Byte data;
        int weight;//结点权值
        Node left;
        Node right;
        public Node(Byte data, int weight) {
            this.data = data;
            this.weight = weight;
        }

        @Override
        public int compareTo(Node o) {
            //表示从小到大排序
            return this.weight - o.weight;
        }

        @Override
        public String toString() {
            return "Node{" + "data=" + data + ", weight=" + weight + '}';
        }
    }
}